package main

import (
	"bufio"
	"fmt"
	"hash/crc64"
	"io"
	"math"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	log "github.com/happyh/go-logging"
	"github.com/happyh/pflag"
	"github.com/nightlyone/lockfile"
)

type downloadInfo struct {
	Url     string
	Headers map[string]string
}
type Task struct {
	Index   int
	Percent int
	Speed   int
	Err     error
}

var progressTick = 2 //每个任务2秒刷新一次显示进度

func main() {
	//使用cobra进行解析？
	var parallel int
	var cookie string
	var outputfilename string
	var referer string
	var taskfile string
	var noScreen bool
	pflag.IntVarP(&parallel, "parallel", "p", 10, "并发的协程数")
	pflag.StringVarP(&cookie, "cookie", "c", "", "cookie")
	pflag.StringVarP(&outputfilename, "out", "o", "", "保存文件名")
	pflag.StringVarP(&referer, "referer", "r", "", "referer 地址")
	pflag.BoolVarP(&noScreen, "noscreen", "s", false, "是否关闭screen模式，默认是screen模式")
	pflag.StringVarP(&taskfile, "taskfile", "t", "", "下载参数文件，如果指定则将文件的第一行读取解析为参数，目的是为了避免在命令行显示下载的url")
	// 设置自定义用法信息函数
	pflag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		pflag.PrintDefaults()
	}

	pflag.Parse()

	// 判断是否没有传递任何参数
	if len(os.Args) == 1 {
		pflag.Usage()
		os.Exit(0)
	}

	switch runtime.GOOS {
	case "windows":
		noScreen = true
	}

	if noScreen && taskfile != "" {
		commandcotent, _ := readFirstLine(taskfile)
		commands := strings.Split(commandcotent, " ")
		pflag.ParseContent(commands)
		noScreen = true
	}

	logfilename := os.Args[0] + ".log"
	log.Init(logfilename, 6)

	positionalArgs := pflag.Args()
	if noScreen {
		var wg sync.WaitGroup
		progressTick = 2 * int(math.Min(float64(len(positionalArgs)), 8))
		concurrencyLimit := make(chan struct{}, 8)
		for _, taskurl := range positionalArgs {
			lowerurl := strings.ToLower(taskurl)
			if !strings.HasPrefix(lowerurl, "http://") && !strings.HasPrefix(lowerurl, "https://") {
				log.Log().Infof("下载url：%s 失败， 不是以http或者https开头\n", taskurl)
				continue
			}
			wg.Add(1)
			// 向通道发送一个空结构体，表示一个新的任务开始
			concurrencyLimit <- struct{}{}
			log.Log().Infof("下载url：%s \n", taskurl)
			go func(url string) {
				defer wg.Done()
				defer func() { <-concurrencyLimit }() // 任务完成后释放一个通道空位
				dowdownloadTask(url, parallel, cookie, referer, outputfilename)
			}(taskurl)
			time.Sleep(2 * time.Second)
		}
		wg.Wait()
	} else {
		// 构造要执行的命令
		command := fmt.Sprintf("%s --noscreen -p %d", os.Args[0], parallel)
		if referer != "" {
			command = command + " -r '" + referer + "'"
		}
		if cookie != "" {
			command = command + " -c '" + cookie + "'"
		}
		if outputfilename != "" {
			command = command + " -o '" + outputfilename + "'"
		}
		if taskfile != "" {
			command = command + " -t '" + taskfile + "'"
		}

		for _, taskurl := range positionalArgs {
			lowerurl := strings.ToLower(taskurl)
			if !strings.HasPrefix(lowerurl, "http://") && !strings.HasPrefix(lowerurl, "https://") {
				log.Log().Infof("下载url：%s 失败， 不是以http或者https开头\n", taskurl)
				continue
			}
			command = command + " '" + taskurl + "'"
		}
		pwd, _ := os.Getwd()
		log.Log().Infof("当前目录：%s, screen 后台执行任务%s\n", pwd, command)

		// 使用screen执行命令
		cmd := exec.Command("screen", "-dmS", "my_screen", "bash", "-c", command)
		//cmd := exec.Command("/data/public/screen-4.8.0/screen", "-dmS", "my_screen", "bash", "-c", command)
		if err := cmd.Start(); err != nil {
			log.Log().Error("执行screen命令时出错:", err)
			os.Exit(1)
		} else {
			pid := cmd.Process.Pid
			log.Log().Info("执行screen命令成功:", cmd, ",pid:", pid)
		}
	}
}
func readFirstLine(filename string) (string, error) {
	file, err := os.Open(filename)
	if err != nil {
		log.Log().Infof("打开文件：%s 失败: %v", filename, err)
		return "", err
	}
	defer file.Close()

	reader := bufio.NewReader(file)
	line, err := reader.ReadString('\n')
	if err != nil && err != io.EOF {
		log.Log().Infof("读取文件：%s 失败: %v", filename, err)
		return "", err
	}

	// 去除行尾的换行符和空白字符
	line = strings.TrimSpace(line)
	if line == "" {
		log.Log().Infof("文件：%s 为空", filename)
		return "", fmt.Errorf("file is empty")
	}

	return line, nil
}

func dowdownloadTask(url string, parallel int, cookie, referer, outputfilename string) {
	// Check if the filePath is a URL.
	lowerurl := strings.ToLower(url)
	if strings.HasPrefix(lowerurl, "http://") || strings.HasPrefix(lowerurl, "https://") {
		info := downloadInfo{
			Url:     strings.TrimSpace(url),
			Headers: map[string]string{"User-Agent": "Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/53.0.2763.0 Safari/537.36"},
		}

		if cookie != "" {
			info.Headers["Cookie"] = cookie
		}
		if referer != "" {
			info.Headers["Referer"] = referer
		}

		downloadResource(info, parallel, outputfilename)
	}
}

func downloadResource(info downloadInfo, parallel int, outfilename string) {
	filename, filesize, crc64, err := downloadHeader(info)
	if err != nil {
		log.Log().Error("获取文件头信息失败:", err)
		return
	}

	if outfilename != "" {
		filename = outfilename
	}

	fileInfo, err := os.Stat(filename)
	if err == nil && filesize == fileInfo.Size() {
		log.Log().Info("文件已经存在:", filename)
		return
	} else if err == nil && filesize != fileInfo.Size() {
		log.Log().Info("文件已经存在:", filename, ",但是文件大小不对，重新下载")
	}

	pwd, _ := os.Getwd()
	// 创建一个新的锁文件实例
	lf, err := lockfile.New(pwd + "/" + filename + ".lock")
	if err != nil {
		log.Log().Errorf("创建锁文件时出错: %v", err)
		return
	}

	err = lf.TryLock()
	if err = lf.TryLock(); err != nil {
		log.Log().Errorf("Cannot lock %q, reason: %v", lf, err)
		return
	}
	defer lf.Unlock()

	// 创建一个临时目录来保存分片
	tmpdir := pwd + "/" + filename + "downloading"
	err = os.Mkdir(tmpdir, os.ModePerm)
	if err != nil && !os.IsExist(err) {
		log.Log().Error("创建临时目录失败:", err)
		return
	}

	if filesize > 10*1024*1024*1024 && parallel == 10 {
		parallel = int(filesize/(1024*1024*1024)) + 1
	}
	if filesize < 1024*1024 && parallel == 10 {
		parallel = int(filesize/(100*1024)) + 1
	}

	part_size := (filesize-1)/int64(parallel) + 1

	task_info := make(chan Task)
	defer close(task_info)

	partfilename := make([]string, parallel)
	var wg sync.WaitGroup
	for i := 0; i < parallel; i++ {
		range_begin := int64(i) * part_size
		range_end := int64(i+1)*part_size - 1
		if range_end >= filesize {
			range_end = filesize - 1
		}
		wg.Add(1)
		go func(index int, range_begin, range_end int64) {
			defer wg.Done()
			fullfileanme := tmpdir + "/" + filename + "." + strconv.FormatInt(range_begin, 10) + "-" + strconv.FormatInt(range_end, 10)
			partfilename[index] = fullfileanme
			tryMax := 10
			tryCount := 0
			var err error
			for tryCount < tryMax {
				tryCount += 1
				err = downloadPart(index, info, range_begin, range_end, fullfileanme, filename, task_info)
				if err != nil {
					log.Log().Errorf("%s 第 %d 次下载分片%d失败: %v\n", filename, tryCount, index, err)
					time.Sleep(30 * time.Second)
				} else {
					log.Log().Infof("%s 分片%d下载完成。\n", filename, index)
					break
				}
			}

			var task Task
			task.Index = index
			task.Percent = 100
			task.Speed = 0
			task.Err = err
			task_info <- task
		}(i, range_begin, range_end)
	}

	tick := time.Tick(time.Duration(progressTick) * time.Second)
	arTasks := make([]Task, parallel)
	isError := false
	for {
		finish := 0
		select {
		case p := <-task_info:
			arTasks[p.Index] = p
		case <-tick:
			var Percent = 0.0
			var Speed = 0
			var strSpeed = " "
			for _, v := range arTasks {
				if v.Percent >= 100 {
					finish++
				}
				if v.Err != nil {
					isError = true
					finish++
				}
				Percent += float64(v.Percent)
				Speed += v.Speed
				//			strSpeed += humanSize(int64(v.Speed)) + " "
			}
			Percent = Percent / float64(parallel)
			leftTime := calculateLeftTime(int64(Speed), filesize, int64(float64(filesize)*Percent/100))
			fmt.Printf("\r%s %d", filename+": "+strconv.FormatFloat(Percent, 'f', 2, 64)+"% "+humanSize(int64(Speed))+" "+leftTime+strSpeed, finish)
			fmt.Printf("\033[K")
		}
		if finish >= parallel {
			break
		}
	}

	wg.Wait()

	if !isError {
		err = mergedFiles(partfilename, pwd+"/"+filename, crc64)
		if err == nil {
			os.RemoveAll(tmpdir)
		} else {
			log.Log().Error("合并文件失败，err:", err)
		}
	}
}

func downloadHeader(info downloadInfo) (filename string, filesize int64, crc64 uint64, err error) {
	req, err := http.NewRequest("GET", info.Url, nil)
	if err != nil {
		return "", 0, 0, err
	}
	req.Header.Set("Range", "bytes=0-0")
	for key, value := range info.Headers {
		req.Header.Set(key, (value))
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return "", 0, 0, err
	}
	defer resp.Body.Close()

	// 检查响应状态码，确保是成功或部分内容（206）
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
		return "", 0, 0, fmt.Errorf("请求文件名文件大小等信息失败，状态码:%d,url:%s", resp.StatusCode, info.Url)
	}

	re := regexp.MustCompile(`filename\*=UTF-8\'\'(.+)`)
	match := re.FindStringSubmatch(resp.Header.Get("Content-Disposition"))
	if len(match) < 1 {
		parsedURL, err := url.Parse(info.Url)
		if err != nil {
			return "", 0, 0, fmt.Errorf("获取url文件名失败")
		}
		filename, err = url.QueryUnescape(filepath.Base(parsedURL.Path))
		if err != nil {
			return "", 0, 0, fmt.Errorf("Error decoding filename: %v", err)
		}
	} else {
		filename, err = url.QueryUnescape(match[1])
		if err != nil {
			return "", 0, 0, fmt.Errorf("Error decoding filename: %v", err)
		}
	}
	filename = sanitizeFilename(filename)

	match = strings.Split(resp.Header.Get("Content-Range"), "/") //bytes 0-0/1256
	if len(match) != 2 {
		return "", 0, 0, fmt.Errorf("Content-Range格式错误:%v", resp.Header.Get("Content-Range"))
	}

	filesize, err = strconv.ParseInt(match[1], 10, 64)
	if err != nil {
		return "", 0, 0, fmt.Errorf("解析文件大小错误 %v", match[2])
	}

	crc64, err = strconv.ParseUint(strings.TrimSpace(resp.Header.Get("x-oss-hash-crc64ecma")), 10, 64)

	return filename, filesize, crc64, nil
}

func downloadPart(index int, info downloadInfo, rangeBegin int64, rangeEnd int64, fullFilename string, filename string, p chan<- Task) error {
	// 创建HTTP请求并处理错误
	req, err := http.NewRequest("GET", strings.TrimSpace(info.Url), nil)
	if err != nil {
		return fmt.Errorf("创建HTTP请求失败: %w", err)
	}

	// 获取已下载文件大小，妥善处理文件不存在的情况
	var fileSize int64
	fileInfo, err := os.Stat(fullFilename)
	if err != nil {
		if !os.IsNotExist(err) { // 仅忽略"文件不存在"错误，其他错误需处理
			return fmt.Errorf("获取文件信息失败: %w", err)
		}
		fileSize = 0 // 文件不存在时视为0字节
	} else {
		fileSize = fileInfo.Size()
	}

	rangeSize := rangeEnd - rangeBegin + 1

	// 检查文件状态并处理
	switch {
	case fileSize == rangeSize:
		log.Log().Info(fmt.Sprintf("%s %d 已经下载完成, filesize:%d, rangesize:%d",
			filename, index, fileSize, rangeSize))
		return nil

	case fileSize > rangeSize:
		return fmt.Errorf("%s %d 下载错误: 文件大小(%d)大于分片大小(%d)",
			filename, index, fileSize, rangeSize)

	case fileSize > 0 && fileSize < rangeSize:
		log.Log().Info(fmt.Sprintf("%s %d 继续下载, filesize:%d, rangesize:%d",
			filename, index, fileSize, rangeSize))

	case fileSize == 0:
		log.Log().Info(fmt.Sprintf("%s %d 开始下载, rangesize:%d",
			filename, index, rangeSize))
	}

	// 设置请求头
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", fileSize+rangeBegin, rangeEnd))
	for key, value := range info.Headers {
		req.Header.Set(key, strings.TrimSpace(value))
	}

	// 发送HTTP请求
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("发送请求失败: %w", err)
	}
	defer resp.Body.Close() // 确保响应体最终关闭

	// 检查响应状态码
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
		return fmt.Errorf("请求分片失败，状态码：%d", resp.StatusCode)
	}

	// 打开文件准备写入
	file, err := os.OpenFile(fullFilename, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		return fmt.Errorf("打开文件失败: %w", err)
	}
	defer file.Close() // 确保文件最终关闭

	// 初始化任务状态
	task := Task{
		Index:   index,
		Err:     nil,
		Speed:   0,
		Percent: 0,
	}

	// 配置缓冲区和IO工具
	const bufferSize = 10 * 1024 // 常量定义更规范
	reader := bufio.NewReaderSize(resp.Body, bufferSize)
	writer := bufio.NewWriter(file)
	defer writer.Flush() // 确保缓冲区数据最终写入文件

	// 进度跟踪变量
	var (
		written     int64
		lastReceive int64
		lastTime    = time.Now()
		interval    = 10 * time.Second
		totalSize   = rangeSize
	)

	// 读写数据循环
	buffer := make([]byte, bufferSize)
	for {
		nr, er := reader.Read(buffer)
		if nr > 0 {
			nw, ew := writer.Write(buffer[:nr])
			if nw > 0 {
				written += int64(nw)
				lastReceive += int64(nw)

				// 定期更新进度
				if time.Since(lastTime) > interval {
					task.Percent = int((written + fileSize) * 100 / totalSize)
					task.Speed = int(lastReceive / int64(time.Since(lastTime).Seconds()))
					p <- task

					// 重置进度跟踪
					lastTime = time.Now()
					lastReceive = 0
				}
			}

			// 处理写入错误
			if ew != nil {
				return fmt.Errorf("写入文件失败: %w", ew)
			}
			if nr != nw {
				return io.ErrShortWrite
			}
		}

		// 处理读取错误
		if er != nil {
			if er != io.EOF {
				return fmt.Errorf("读取响应数据失败: %w", er)
			}
			break // 正常结束循环
		}
	}

	// 发送最终完成状态
	task.Percent = 100
	task.Speed = 0
	p <- task

	return nil
}

func calculateLeftTime(speed int64, contentSize int64, downloadedSize int64) string {
	if speed == 0 {
		return "NaN"
	}

	remainingBytes := contentSize - downloadedSize
	remainingSeconds := int64(float64(remainingBytes) / float64(speed))
	duration := time.Duration(remainingSeconds) * time.Second
	return duration.String()
}

func mergedFiles(files []string, outputFile string, crc64 uint64) error {
	outFile, err := os.Create(outputFile)
	if err != nil {
		return fmt.Errorf("Error creating output file:%v", err)
	}
	defer outFile.Close()

	// 逐个打开文件并合并内容
	for _, filename := range files {
		// 打开文件
		inFile, err := os.Open(filename)
		if err != nil {
			return fmt.Errorf("Error opening file %s: %v\n", filename, err)
		}
		defer inFile.Close()

		// 复制文件内容到输出文件
		_, err = io.Copy(outFile, inFile)
		if err != nil {
			return fmt.Errorf("Error writing to output file:%v", err)
		}
	}

	if crc64 != 0 {
		crc64check, err := crc64Checksum(outputFile)
		if err != nil {
			return err
		}

		if crc64check != crc64 {
			return fmt.Errorf("crc64Checksum failed,crc64check:%v,crc64:%v", crc64check, crc64)
		}
	}

	return nil
}

func humanSize(sizeBytes int64) string {
	if sizeBytes == 0 {
		return "0B"
	}

	const unit = 1024
	units := []string{"B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}
	i := int(math.Floor(math.Log(float64(sizeBytes)) / math.Log(float64(unit))))
	p := int64(math.Pow(float64(unit), float64(i)))
	s := float64(sizeBytes) / float64(p)
	// 格式化到小数点后两位
	s = math.Round(s*100) / 100

	return fmt.Sprintf("%.2f%s", s, units[i])
}

func sanitizeFilename(filename string) string {
	// 非法字符列表
	illegalChars := "/\\:*?\"<>|"

	// 遍历非法字符列表，并使用strings.ReplaceAll替换它们
	for _, char := range illegalChars {
		filename = strings.ReplaceAll(filename, string(char), "_")
	}

	return filename
}

// 计算文件的CRC64校验和
func crc64Checksum(filePath string) (uint64, error) {
	// 打开文件
	file, err := os.Open(filePath)
	if err != nil {
		return 0, err
	}
	defer file.Close()

	hasher := crc64.New(crc64.MakeTable(crc64.ECMA))

	// 读取文件内容并更新哈希
	if _, err := io.Copy(hasher, file); err != nil {
		return 0, err
	}

	// 获取最终的CRC64校验和
	return hasher.Sum64(), nil
}
